`timescale 1ns / 1ps

module aru_reduce_stage1_tb;

    // 时钟和复位
    logic clk;
    logic rst_n;

    // 参数定义
    parameter int P_ARU = 4;
    parameter int N0 = 8;
    parameter int DATA_WIDTH = P_ARU * N0;  // 32

    // 时钟生成
    initial begin
        clk = 0;
        forever #5 clk = ~clk;  // 10ns周期，100MHz
    end

    // 接口实例化
    aru_reduce_ctrl_if u_ctrl_if (
        clk,
        rst_n
    );
    aru_payload_if u_payload_if (
        clk,
        rst_n
    );  // 输入到transpose
    aru_reduce_pld_if u_transpose_if (
        clk,
        rst_n
    );  // transpose输出到stage1
    aru_reduce_pld_if u_stage1_if (
        clk,
        rst_n
    );  // stage1输出到stage2
    aru_reduce_pld_if u_stage2_if (
        clk,
        rst_n
    );  // stage2输出到stage3
    aru_reduce_pld_if u_stage3_if (
        clk,
        rst_n
    );  // stage3输出（最终输出）

    // DUT实例化 - Transpose模块
    aru_reduce_transpose u_transpose (
        .clk                (clk),
        .rst_n              (rst_n),
        .u_aru_ctrl_if      (u_ctrl_if.in),
        .u_aru_payload_if   (u_payload_if.in),
        .u_aru_reduce_pld_if(u_transpose_if.out)
    );

    // DUT实例化 - Stage1模块
    aru_reduce_stage1 u_stage1 (
        .clk                      (clk),
        .rst_n                    (rst_n),
        .u_aru_ctrl_if            (u_ctrl_if.in),
        .u_aru_reduce_transpose_if(u_transpose_if.in),
        .u_aru_reduce_stage1_if   (u_stage1_if.out)
    );

    // DUT实例化 - Stage2模块
    aru_reduce_stage2 u_stage2 (
        .clk                   (clk),
        .rst_n                 (rst_n),
        .u_aru_ctrl_if         (u_ctrl_if.in),
        .u_aru_reduce_stage1_if(u_stage1_if.in),
        .u_aru_reduce_stage2_if(u_stage2_if.out)
    );

    // DUT实例化 - Stage3模块
    aru_reduce_stage3 u_stage3 (
        .clk                   (clk),
        .rst_n                 (rst_n),
        .u_aru_ctrl_if         (u_ctrl_if.in),
        .u_aru_reduce_stage2_if(u_stage2_if.in),
        .u_aru_reduce_stage3_if(u_stage3_if.out)
    );

    // 测试任务和函数

    // 复位任务
    task reset();
        rst_n               = 0;
        u_ctrl_if.vld       = 0;
        u_ctrl_if.reduce_m  = 0;
        u_ctrl_if.reduce_n  = 0;
        u_ctrl_if.reduce_op = 0;
        u_ctrl_if.slice_m   = 0;
        u_ctrl_if.slice_n   = 0;
        u_payload_if.vld    = 0;
        u_payload_if.dat    = '0;
        u_payload_if.sdb    = '0;
        u_stage3_if.rdy     = 1;  // 改为stage3的ready
        repeat (10) @(posedge clk);
        rst_n = 1;
        repeat (5) @(posedge clk);
        $display("[%0t] Reset completed", $time);
    endtask

    // 配置任务
    task send_config(input logic reduce_m, input logic reduce_n, input logic [1:0] reduce_op, input int slice_m,
                     input int slice_n);
        @(posedge clk);
        u_ctrl_if.vld       = 1;
        u_ctrl_if.reduce_m  = reduce_m;
        u_ctrl_if.reduce_n  = reduce_n;
        u_ctrl_if.reduce_op = reduce_op;
        u_ctrl_if.slice_m   = slice_m;
        u_ctrl_if.slice_n   = slice_n;
        @(posedge clk);
        wait (u_ctrl_if.rdy);
        @(posedge clk);
        u_ctrl_if.vld = 0;
        $display("[%0t] Config sent: reduce_m=%0d, reduce_n=%0d, reduce_op=%0d, slice_m=%0d, slice_n=%0d", $time,
                 reduce_m, reduce_n, reduce_op, slice_m, slice_n);
    endtask

    // 发送多周期数据任务
    task send_data_multi_cycles(input int start_value);
        int current_value;
        current_value = start_value;

        for (int cyc = 0; cyc < 5; cyc++) begin
            @(posedge clk);
            u_payload_if.vld = 1;

            // 最后一个周期：设置eom和eon
            if (cyc == 4) begin
                u_payload_if.sdb.eom   = 1;
                u_payload_if.sdb.eon   = 1;
                u_payload_if.sdb.vld_m = 4;
                u_payload_if.sdb.vld_n = 4;
            end else begin
                u_payload_if.sdb.eom   = 0;
                u_payload_if.sdb.eon   = 0;
                u_payload_if.sdb.vld_m = 4;
                u_payload_if.sdb.vld_n = 8;
            end

            // 填充递增整数数据（以BF16格式表示整数）
            for (int i = 0; i < DATA_WIDTH; i++) begin
                u_payload_if.dat.dat[i].sign = 1'b0;  // 正数
                u_payload_if.dat.dat[i].exp  = 8'h80;  // 固定指数
                u_payload_if.dat.dat[i].mant = current_value[6:0];  // 取低7位作为尾数
                current_value                = current_value + 1;
            end

            $display("[%0t] Input Cycle %0d/5: Data sent (values: %0d-%0d), eom=%0d, eon=%0d", $time, cyc + 1,
                     current_value - DATA_WIDTH, current_value - 1, u_payload_if.sdb.eom, u_payload_if.sdb.eon);

            // 等待握手
            while (!u_payload_if.rdy) @(posedge clk);
        end

        @(posedge clk);
        u_payload_if.vld       = 0;
        u_payload_if.sdb.eom   = 0;
        u_payload_if.sdb.eon   = 0;
        u_payload_if.sdb.vld_m = 0;
        u_payload_if.sdb.vld_n = 0;
        $display("[%0t] 5-cycle transmission completed (data range: %0d - %0d)", $time, start_value, current_value - 1);
    endtask

    // 等待输出任务
    task wait_output(input int num_cycles);
        repeat (num_cycles) @(posedge clk);
    endtask

    // 测试场景

    // 测试1: MAX操作
    task test_max_operation();
        $display("\n========== Test 1: MAX Operation ==========");
        send_config(.reduce_m(1), .reduce_n(0), .reduce_op(2'b00), .slice_m(4), .slice_n(36));
        send_data_multi_cycles(.start_value(1));
        wait_output(25);
        $display("Test 1 completed\n");
    endtask

    // 测试2: MIN操作
    task test_min_operation();
        $display("\n========== Test 2: MIN Operation ==========");
        send_config(.reduce_m(1), .reduce_n(0), .reduce_op(2'b01), .slice_m(4), .slice_n(36));
        send_data_multi_cycles(.start_value(1));
        wait_output(25);
        $display("Test 2 completed\n");
    endtask

    // 测试3: SUM操作
    task test_sum_operation();
        $display("\n========== Test 3: SUM Operation ==========");
        send_config(.reduce_m(1), .reduce_n(0), .reduce_op(2'b10), .slice_m(4), .slice_n(36));
        send_data_multi_cycles(.start_value(1));
        wait_output(25);
        $display("Test 3 completed\n");
    endtask

    // 测试4: AVG操作
    task test_avg_operation();
        $display("\n========== Test 4: AVG Operation ==========");
        send_config(.reduce_m(1), .reduce_n(0), .reduce_op(2'b11), .slice_m(4), .slice_n(36));
        send_data_multi_cycles(.start_value(1));
        wait_output(25);
        $display("Test 4 completed\n");
    endtask

    // 测试5: BYPASS模式
    task test_bypass_mode();
        $display("\n========== Test 5: BYPASS Mode ==========");
        send_config(.reduce_m(0), .reduce_n(0), .reduce_op(2'b00), .slice_m(4), .slice_n(36));
        send_data_multi_cycles(.start_value(1));
        wait_output(25);
        $display("Test 5 completed\n");
    endtask

    // 测试6: 多次连续操作
    task test_continuous_operations();
        $display("\n========== Test 6: Multiple Continuous Operations ==========");
        send_config(.reduce_m(1), .reduce_n(0), .reduce_op(2'b00), .slice_m(4), .slice_n(36));
        send_data_multi_cycles(.start_value(1));
        wait_output(25);

        send_config(.reduce_m(1), .reduce_n(0), .reduce_op(2'b10), .slice_m(4), .slice_n(36));
        send_data_multi_cycles(.start_value(1));
        wait_output(25);
        $display("Test 6 completed\n");
    endtask

    // 主测试流程
    initial begin
        $display("========== Simulation Start ==========");
        $display("Testing: Full Reduce Pipeline");
        $display("  Transpose -> Stage1 -> Stage2 -> Stage3");
        $display("Data Send Rule: 5 cycles per test, last cycle with eom/eon=1");
        $display("Data Value Rule: Starting from 1.0, auto-increment by 1.0\n");

        reset();

        test_max_operation();
        test_min_operation();
        test_sum_operation();
        test_avg_operation();
        test_bypass_mode();
        test_continuous_operations();

        repeat (50) @(posedge clk);

        $display("========== Simulation End ==========");
        $finish;
    end

    // 超时保护
    initial begin
        #250000;  // 增加超时时间，因为多了两级stage
        $display("ERROR: Simulation timeout!");
        $finish;
    end

    // 波形dump
    initial begin
        $fsdbDumpfile("./aru_reduce_stage1.fsdb");
        $fsdbDumpvars("+all");
    end

endmodule
